# -*- coding:utf-8 -*-#
# @Time:2023/6/26 13:35
# @Author:Adong
# @Software:PyCharm

"""
基于幅值与相角波动性的变压器风机噪声分离步骤demo
"""


import librosa
import numpy as np
import matplotlib.pyplot as plt
from scipy.fftpack import fft, ifft

class AAF:
    def __init__(self,y,sr):
        self.sr = 48000
        self.y = librosa.resample(y=y, orig_sr=sr, target_sr=sr)[:sr]
    def amplitude(self):
        n_fft = int(0.2 * self.sr)  # 帧长度是0.2s
        hop_length = int(0.02 * self.sr)    # 帧移0.02s
        noise_stft = librosa.stft(y,n_fft = n_fft,hop_length=hop_length,window="hamm")
        amp = noise_stft.real
        amp_var = []            # 所有频率幅值的方差
        for hz in amp:
            amp_var.append(np.var(hz))
        # plt.figure()
        # plt.plot(amp_var)
        # plt.show(block=True)

        '''计算50hz倍频点附近频点的方差值平均'''
        hz_50 = np.arange(50,len(noise_stft),50)        # 找出50hz倍频
        self.hz_50 = hz_50
        W_win = 10
        alphaX = []         # 保存50hz倍频附近W_win个频率点方差的平均值
        Ptrans = []             # 变压器频率
        Pfan = []               # 风扇频率
        for idx,x in enumerate(hz_50):
            piece = amp_var[x-W_win:x+W_win]        # 找出50hz倍频附近W_win个频率点的方差
            mean = np.mean(piece)/2                 # 计算平均值
            alphaX.append(mean)
            if mean > amp_var[x]:           # 如果平均值大于当前50倍频，则认为是本体
                Ptrans.append(x)
            else:
                Pfan.append(x)

        # y_fft = fft(y)  # 对y做傅里叶变换
        # y_fft_new = []
        # for idx,x in enumerate(y_fft):
        #     if idx in Pfan:
        #         y_fft_new.append(0)
        #     elif idx in Ptrans:
        #         y_fft_new.append(x)
        #     else:
        #         y_fft_new.append(x)
        # y_ifft = ifft(y_fft_new)  # 逆傅里叶变换
        # plt.figure()
        # plt.plot(y_ifft)
        # plt.show(block=True)
        # from data_pre_processing import Universal_tool
        # Universal_tool.array2wav(y_ifft, 'AAF.wav', self.sr)
        return Ptrans,Pfan

    def if_position_phase_fluctuations(self,hz):
        """
        判断相角波动超90度的次数是否超过阈值
        :param hz:
        :return:
        """
        k = hz
        n_fft = int(10/k * self.sr)
        hop_length = int(1/k * self.sr)
        noise_stft = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, window="hamm")
        phase = noise_stft.imag[hz-1]
        amp = noise_stft.real[hz-1]
        mo = (amp**2 + phase**2)**0.5
        phase_1 = phase / mo                             # 模值除以本身等比缩减至1，相位随之改变
        amp_1 = amp / mo
        # plt.figure()
        # plt.plot(amp_1[:100],phase_1[:100])
        # plt.show(block=True)
        '''计算相角'''
        xiangjiao = []
        up_90 = []
        down_90 = []
        for idx in range(len(amp_1)-1):
            a = amp_1[idx]
            a1 = amp_1[idx+1]
            b = phase_1[idx]
            b1 = phase_1[idx+1]
            xj = ((a-a1)**2+(b-b1)**2)*0.5
            xiangjiao.append(xj)
            if xj > 2**0.5:
                up_90.append(xj)
            else:
                down_90.append(xj)
        miu = 0.05
        if len(up_90) > miu * len(xiangjiao):
            return 0
        else:
            return 1

    def phase_fluctuations(self):
        hz_50 = self.hz_50
        Ptrans = []
        Pfan = []
        for i in hz_50:
            if self.if_position_phase_fluctuations(i) == 0:
                Pfan.append(i)
            else:
                Ptrans.append(i)
        return  Ptrans,Pfan





if __name__ == '__main__':
    filepath = './data/wav_data_V4/normal/环境声响_正常运行_normal_01.wav'
    y, sr = librosa.load(filepath)
    test = AAF(y,sr)
    Ptrans1, Pfan1 = test.amplitude()
    # Ptrans2, Pfan2 = test.phase_fluctuations()